import requests
import json
import os
import re
from urllib.parse import urlparse
import time
import sys
import io
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
from threading import Lock  

# 全局配置参数
KNOWLEDGE_BASE_ID = "7308316258628471"
FOLDER_ID = "7308316258628471"
API_VERSION = "17596504052671759196866613"
API_URL = "https://ima.qq.com/cgi-bin/knowledge_tab_reader/get_knowledge_list"
DEFAULT_PAGE_SIZE = 50
REQUEST_DELAY = 1  # 请求间隔延迟（秒）
MAX_WORKERS = 5  # 最大线程数
DOWNLOAD_DELAY = 0.5  # 下载间隔延迟（秒）

# 动态请求头参数（需要定期更新）
TRACEPARENT = "00-989dcae1f748d2f19fd7b67be6e5d20a-9f3b3edc8ae9f0f1-01"
X_IMA_BKN = "658036523"
X_IMA_COOKIE = "PLATFORM=H5; CLIENT-TYPE=256053; WEB-VERSION=999.999.999; IMA-GUID=guid-0de94fa12725aa3aecb705500200000a81850a; IMA-Q36=0de94fa12725aa3aecb705500200000a81850a; IMA-IUA=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36 Edg/141.0.0.0; IMA-UID=001a3672d3c04341; IMA-TOKEN=pUXb2+qvvPGVUynh5M8ZkKxkp58kG6H1OLUUGl6Sd+yKNNPOGsVLifVI99c231cgFE6nlnJ2JTbU5oNwsu7GElf2e9JUS65nZJAGYKACAQczKz9GU5Rk66Vg5TORCqRinbH8ZrmaB+289+bNtx7swNKXXIkittmqz1oG4Nm847/eAg==; IMA-REFRESH-TOKEN=EEHw60UM1zAK1q6KnyDg6IfOwdTgGQiSyfjtIbqgeVAvE9xENK/DWc0DcU2eGa0FA/wnaOWY4kHQWjMp0EecmxsE18iLTNAqyQj+S5vTG2WQRH3qO36M51jzKTV6fmL1V9UeY5pgJwdXEAgyXb/0KJaWdtgG+nnsbdbJaf0Du/SX73oBHBlnmQV4GOJlsisl+QxUaQcW; UID-TYPE=2; TOKEN-TYPE=14"


class IMAKnowledgeDownloader:
    def __init__(self):
        self.session = requests.Session()
        self.setup_headers()
        # 设置输出编码
        if sys.platform == "win32":
            sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
            sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')
        
        # 多线程统计变量
        self.success_count = 0
        self.failed_count = 0
        self.lock = Lock()  # 线程锁，用于保护共享变量
        
    def setup_headers(self):
        """设置请求头"""
        self.headers = {
            'Accept-Encoding': 'gzip, deflate, br, zstd',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6',
            'Cache-Control': 'no-cache',
            'Connection': 'keep-alive',
            'Cookie': 'ptcz=118bc53b1e2f89b2c5d6ad8a72460ce1ced3243d37a6b9242e1d106025a36e4b; pgv_pvid=6803744288; _qimei_h38=0de94fa12725aa3aecb705500200000a81850a; _qimei_q32=a2b406e509133fea499cae10de653004; _qimei_q36=21e3402eff2542d815f21b2a30001fd18508; o_cookie=452438861; yyb_muid=0CCBE4111A5D63A80DDFF0681B2762C4; _qimei_uuid42=1970c0d1e24100345053560bce0008f1173b91dba9; RK=TdeMUGxOSy; _qimei_i_3=6ffe2dd09652578e9595f7385b8027b3f2eef2f3170f0a85bd8d7e0c2e9a2665623033943a89e2a7b6ae; _qimei_fingerprint=07e3e8ae92d50da0edd501a5d2f0edb1; _qimei_i_2=79b97680940f53dfc195ab660d8676e4f5baacf0170d0183e68d795b2693206d3737369d3d89eaa0beb4; _qimei_i_1=7ec54582c10f59df97c6ff325c8676b1a6edf5a441090482b18a28582493206c6163669d3981e7dc8786c0fa; omgid=0_25715e9fc2a1c; pac_uid=0_XZ1Zc2y04xnfs',
            'Host': 'ima.qq.com',
            'Origin': 'https://ima.qq.com',
            'Pragma': 'no-cache',
            'Referer': 'https://ima.qq.com/wikis',
            'Sec-Fetch-Dest': 'empty',
            'Sec-Fetch-Mode': 'cors',
            'Sec-Fetch-Site': 'same-origin',
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36 Edg/141.0.0.0',
            'accept': 'application/json',
            'content-type': 'application/json',
            'extension_version': '999.999.999',
            'from_browser_ima': '1',
            'sec-ch-ua': '"Microsoft Edge";v="141", "Not?A_Brand";v="8", "Chromium";v="141"',
            'sec-ch-ua-mobile': '?0',
            'sec-ch-ua-platform': '"Windows"',
            'traceparent': TRACEPARENT,
            'x-ima-bkn': X_IMA_BKN,
            'x-ima-cookie': X_IMA_COOKIE
        }
        
    def get_knowledge_list(self, cursor="", limit=DEFAULT_PAGE_SIZE):
        """
        获取知识列表数据
        
        Args:
            cursor: 分页游标，第一页为空字符串
            limit: 每页数量
            
        Returns:
            dict: 返回的JSON数据
        """
        # 使用固定的API端点
        url = API_URL
        
        # 请求数据
        request_data = {
            "sort_type": 0,
            "need_default_cover": True,
            "knowledge_base_id": KNOWLEDGE_BASE_ID,
            "folder_id": FOLDER_ID,
            "cursor": cursor,
            "limit": limit,
            "version": API_VERSION,
            "ext_info": {}
        }
        
        try:
            print(f"正在请求URL: {url}")
            print(f"请求数据: {json.dumps(request_data, ensure_ascii=False, indent=2)}")
            
            response = self.session.post(
                url, 
                headers=self.headers, 
                json=request_data,
                timeout=30
            )
            response.raise_for_status()
            
            result = response.json()
            print(f"响应状态码: {response.status_code}")
            
            # 检查响应是否成功
            if result.get('code') == 0 or result.get('msg') == 'ok':
                knowledge_list = result.get('knowledge_list', [])
                print(f"成功获取知识列表，共 {len(knowledge_list)} 条记录")
                return result
            else:
                print(f"API返回错误代码: {result.get('code')}, 消息: {result.get('msg')}")
                return None
                
        except requests.exceptions.RequestException as e:
            print(f"请求失败: {e}")
            return None
        except json.JSONDecodeError as e:
            print(f"JSON解析失败: {e}")
            return None
        except Exception as e:
            print(f"其他错误: {e}")
            return None
    
    def get_all_knowledge_list(self):
        """
        获取所有页面的知识列表数据
        
        Returns:
            dict: 合并后的完整数据
        """
        all_knowledge_list = []
        cursor = ""  # 第一页cursor为空
        page_num = 1
        total_size = 0
        
        print("开始分页获取知识列表...")
        
        while True:
            print(f"\n正在获取第 {page_num} 页数据...")
            result = self.get_knowledge_list(cursor=cursor, limit=DEFAULT_PAGE_SIZE)
            
            if not result:
                print(f"第 {page_num} 页获取失败")
                break
            
            # 获取当前页数据
            current_list = result.get('knowledge_list', [])
            if not current_list:
                print("当前页没有数据，分页结束")
                break
            
            all_knowledge_list.extend(current_list)
            
            # 获取下一页的cursor
            next_cursor = result.get('next_cursor', '')
            is_end = result.get('is_end', False)
            total_size = result.get('total_size', '0')
            
            print(f"第 {page_num} 页获取到 {len(current_list)} 条记录")
            print(f"下一页cursor: {next_cursor}")
            print(f"是否结束: {is_end}")
            
            # 检查是否还有下一页
            if is_end or not next_cursor:
                print("已到达最后一页，分页结束")
                break
            
            cursor = next_cursor
            page_num += 1
            
            # 添加延迟避免请求过于频繁
            time.sleep(REQUEST_DELAY)
        
        # 构建完整的返回数据
        if all_knowledge_list:
            complete_data = {
                "code": 0,
                "msg": "ok",
                "knowledge_list": all_knowledge_list,
                "total_size": total_size,
                "is_end": True,
                "next_cursor": ""
            }
            
            # 如果有知识库信息，保留第一个页面的信息
            if page_num > 1:
                first_result = self.get_knowledge_list(cursor="", limit=DEFAULT_PAGE_SIZE)
                if first_result and 'knowledge_base_info' in first_result:
                    complete_data['knowledge_base_info'] = first_result['knowledge_base_info']
            
            print(f"\n分页获取完成！总共获取 {len(all_knowledge_list)} 条记录")
            return complete_data
        else:
            print("未获取到任何数据")
            return None
    
    def sanitize_filename(self, filename):
        """
        清理文件名，移除非法字符
        
        Args:
            filename: 原始文件名
            
        Returns:
            str: 清理后的文件名
        """
        # 移除或替换非法字符
        filename = re.sub(r'[<>:"/\\|?*]', '_', filename)
        # 移除多余的空格和点
        filename = filename.strip('. ')
        # 限制文件名长度
        if len(filename) > 100:
            filename = filename[:100]
        return filename
    
    def fix_duplicate_extension(self, filename):
        """
        修复文件名中的重复后缀
        
        Args:
            filename: 原始文件名
            
        Returns:
            str: 修复后的文件名
        """
        # 获取文件名和扩展名
        name, ext = os.path.splitext(filename)
        
        # 如果扩展名为空，直接返回原文件名
        if not ext:
            return filename
        
        # 检查是否有重复的扩展名
        # 例如：file.pdf.pdf -> file.pdf
        while True:
            new_name, new_ext = os.path.splitext(name)
            if new_ext == ext:
                name = new_name
            else:
                break
        
        # 重新组合文件名
        return name + ext
    
    def is_folder(self, item):
        """
        判断是否为文件夹
        
        Args:
            item: 资源项
            
        Returns:
            bool: 是否为文件夹
        """
        return (item.get('media_type') == 99 and 
                item.get('folder_info') is not None and
                item.get('folder_info', {}).get('folder_id') is not None)
    
    def is_scanned_pdf(self, item):
        """
        判断PDF文件是否为扫描图片格式（不可复制文字）
        
        Args:
            item: 资源项
            
        Returns:
            bool: True表示是扫描图片PDF，False表示是可复制文字的PDF
        """
        # 只处理PDF文件
        if item.get('media_type') != 1:
            return False
        
        # 检查解析状态 - 扫描PDF通常解析进度为0或很低
        parse_progress = item.get('parse_progress', 0)
        if parse_progress == 0:
            return True  # 解析进度为0，很可能是扫描PDF
        
        # 检查摘要状态 - 扫描PDF通常无法生成AI摘要
        summary_state = item.get('summary_state', 0)
        if summary_state == 0:
            return True  # 无法生成摘要，可能是扫描PDF
        
        # 检查是否有摘要内容
        abstract = item.get('abstract', '').strip()
        if not abstract or abstract == '该内容暂未生成摘要':
            return True  # 没有摘要或摘要为空，可能是扫描PDF
        
        # 检查文件大小与解析状态的关系
        file_size = int(item.get('file_size', 0))
        if file_size > 0:
            # 大文件但解析进度低，可能是扫描PDF
            if file_size > 10 * 1024 * 1024 and parse_progress < 50:  # 大于10MB且解析进度低于50%
                return True
        
        # 检查是否有解析错误信息
        parse_err_info = item.get('parse_err_info')
        if parse_err_info is not None:
            return True  # 有解析错误，可能是扫描PDF
        
        # 检查媒体状态
        media_state = item.get('media_state', 0)
        if media_state == 0:
            return True  # 媒体状态为0，可能是扫描PDF
        
        return False  # 通过所有检查，认为是可复制文字的PDF
    
    def should_download_pdf(self, item):
        """
        判断是否应该下载PDF文件（过滤扫描图片PDF）
        
        Args:
            item: 资源项
            
        Returns:
            bool: True表示应该下载，False表示跳过
        """
        # 只处理PDF文件
        if item.get('media_type') != 1:
            return True  # 非PDF文件正常处理
        
        # 检查是否为扫描PDF
        if self.is_scanned_pdf(item):
            return False  # 跳过扫描PDF
        
        return True  # 下载可复制文字的PDF
    
    def get_folder_contents(self, folder_id, folder_name):
        """
        获取文件夹内容
        
        Args:
            folder_id: 文件夹ID
            folder_name: 文件夹名称
            
        Returns:
            dict: 文件夹内容数据
        """
        print(f"正在获取文件夹 '{folder_name}' 的内容...")
        
        # 使用文件夹ID作为新的FOLDER_ID
        original_folder_id = FOLDER_ID
        
        # 创建新的请求数据
        request_data = {
            "sort_type": 0,
            "need_default_cover": True,
            "knowledge_base_id": KNOWLEDGE_BASE_ID,
            "folder_id": folder_id,  # 使用文件夹ID
            "cursor": "",
            "limit": DEFAULT_PAGE_SIZE,
            "version": API_VERSION,
            "ext_info": {}
        }
        
        try:
            response = self.session.post(
                API_URL, 
                headers=self.headers, 
                json=request_data,
                timeout=30
            )
            response.raise_for_status()
            
            result = response.json()
            
            if result.get('code') == 0 or result.get('msg') == 'ok':
                knowledge_list = result.get('knowledge_list', [])
                print(f"成功获取文件夹 '{folder_name}' 内容，共 {len(knowledge_list)} 条记录")
                return result
            else:
                print(f"获取文件夹 '{folder_name}' 内容失败: {result.get('msg')}")
                return None
                
        except Exception as e:
            print(f"获取文件夹 '{folder_name}' 内容时出错: {e}")
            return None
    
    def download_folder_contents(self, folder_data, folder_name, save_dir):
        """
        下载文件夹内容
        
        Args:
            folder_data: 文件夹数据
            folder_name: 文件夹名称
            save_dir: 保存目录
            
        Returns:
            dict: 下载结果统计
        """
        if not folder_data or 'knowledge_list' not in folder_data:
            print(f"文件夹 '{folder_name}' 数据无效")
            return {"success": 0, "failed": 0, "total": 0}
        
        knowledge_list = folder_data.get('knowledge_list', [])
        if not knowledge_list:
            print(f"文件夹 '{folder_name}' 为空")
            return {"success": 0, "failed": 0, "total": 0}
        
        # 获取当前文件夹的实际名称（从current_path或使用传入的folder_name）
        current_dir_name = self.get_current_path_name(folder_data)
        if current_dir_name and current_dir_name != "未知目录":
            actual_folder_name = current_dir_name
        else:
            actual_folder_name = folder_name
        
        # 创建文件夹目录
        folder_save_dir = os.path.join(save_dir, self.sanitize_filename(actual_folder_name))
        os.makedirs(folder_save_dir, exist_ok=True)
        print(f"创建文件夹目录: {os.path.abspath(folder_save_dir)}")
        
        # 重置计数器
        self.success_count = 0
        self.failed_count = 0
        
        # 使用线程池进行多线程下载
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            # 提交所有下载任务
            future_to_item = {}
            for i, item in enumerate(knowledge_list, 1):
                future = executor.submit(
                    self.download_single_document, 
                    item, 
                    folder_save_dir, 
                    i, 
                    len(knowledge_list)
                )
                future_to_item[future] = item
            
            # 处理完成的任务
            completed_count = 0
            for future in as_completed(future_to_item):
                completed_count += 1
                try:
                    result = future.result()
                    if result["success"]:
                        if "folder_result" in result:
                            # 文件夹下载结果
                            folder_result = result["folder_result"]
                            print(f"✓ 文件夹完成: {result['title']} (包含 {folder_result['success']} 个文件)")
                        else:
                            print(f"✓ 文件完成: {result['title']}")
                    else:
                        print(f"✗ 失败: {result['title']} - {result.get('reason', '未知错误')}")
                except Exception as e:
                    print(f"✗ 异常: {future_to_item[future].get('title', '未知')} - {e}")
                    with self.lock:
                        self.failed_count += 1
                
                # 显示进度
                if completed_count % 10 == 0 or completed_count == len(knowledge_list):
                    with self.lock:
                        print(f"文件夹 '{actual_folder_name}' 进度: {completed_count}/{len(knowledge_list)} "
                              f"(成功: {self.success_count}, 失败: {self.failed_count})")
                
                # 添加小延迟避免请求过于频繁
                time.sleep(DOWNLOAD_DELAY)
        
        result = {
            "success": self.success_count,
            "failed": self.failed_count,
            "total": len(knowledge_list)
        }
        
        print(f"文件夹 '{actual_folder_name}' 下载完成! 成功: {self.success_count}, 失败: {self.failed_count}, 总计: {len(knowledge_list)}")
        return result
    
    def download_single_document(self, item, save_dir, index, total):
        """
        下载单个文档或文件夹（用于多线程）
        
        Args:
            item: 文档信息字典
            save_dir: 保存目录
            index: 当前索引
            total: 总数量
            
        Returns:
            dict: 下载结果
        """
        title = item.get('title', f'文档_{index}')
        media_type = item.get('media_type')
        
        # 检查是否为文件夹
        if self.is_folder(item):
            folder_info = item.get('folder_info', {})
            folder_id = folder_info.get('folder_id')
            folder_name = folder_info.get('name', title)
            
            with self.lock:
                print(f"[{index}/{total}] 正在处理文件夹: {folder_name}")
            
            # 获取文件夹内容
            folder_data = self.get_folder_contents(folder_id, folder_name)
            if folder_data:
                # 获取文件夹的实际名称
                actual_folder_name = self.get_current_path_name(folder_data)
                if actual_folder_name == "未知目录":
                    actual_folder_name = folder_name
                
                # 下载文件夹内容
                result = self.download_folder_contents(folder_data, actual_folder_name, save_dir)
                if result['success'] > 0:
                    with self.lock:
                        self.success_count += result['success']
                        self.failed_count += result['failed']
                    return {"success": True, "title": f"文件夹: {actual_folder_name}", "folder_result": result}
                else:
                    with self.lock:
                        self.failed_count += 1
                    return {"success": False, "title": f"文件夹: {actual_folder_name}", "reason": "文件夹下载失败"}
            else:
                with self.lock:
                    self.failed_count += 1
                return {"success": False, "title": f"文件夹: {folder_name}", "reason": "无法获取文件夹内容"}
        
        # 处理普通文件
        jump_url = item.get('jump_url')
        if not jump_url:
            with self.lock:
                self.failed_count += 1
                print(f"[{index}/{total}] 跳过 {title}: 没有jump_url")
            return {"success": False, "title": title, "reason": "没有jump_url"}
        
        # 显示资源类型信息
        media_type_name = "未知类型"
        if media_type == 1:
            media_type_name = "PDF文件"
        elif media_type == 6:
            media_type_name = "公众号文章"
        
        with self.lock:
            print(f"[{index}/{total}] 正在处理: {title} ({media_type_name})")
        
        if self.download_document(jump_url, title, save_dir, media_type):
            with self.lock:
                self.success_count += 1
            return {"success": True, "title": title}
        else:
            with self.lock:
                self.failed_count += 1
            return {"success": False, "title": title, "reason": "下载失败"}
    
    def download_document(self, jump_url, title, save_dir="downloads", media_type=None):
        """
        下载文档内容
        
        Args:
            jump_url: 文档的跳转URL
            title: 文档标题，用于命名文件
            save_dir: 保存目录
            media_type: 媒体类型，6为公众号文章
            
        Returns:
            bool: 下载是否成功
        """
        try:
            # 创建保存目录
            os.makedirs(save_dir, exist_ok=True)
            
            # 清理文件名
            safe_title = self.sanitize_filename(title)
            
            # 先确定文件名，用于检查是否已存在
            if media_type == 6:
                # 公众号文章保存为HTML
                filename = f"{safe_title}.html"
            else:
                # 其他文件类型，先假设为PDF，后续可能会根据实际内容调整
                filename = f"{safe_title}.pdf"
            
            # 修复文件名重复后缀
            filename = self.fix_duplicate_extension(filename)
            
            # 检查文件是否已存在
            file_path = os.path.join(save_dir, filename)
            if os.path.exists(file_path):
                print(f"文件已存在，跳过下载: {filename}")
                return True
            
            # 如果是公众号文章，保存为HTML文件
            if media_type == 6:
                # 文件名已经在前面处理过了
                
                # 创建简化的HTML内容，直接跳转
                html_content = f"""<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <title>{title}</title>
</head>
<body>
    <script>
        // 自动跳转到链接
        window.location.href = "{jump_url}";
    </script>
</body>
</html>"""
                
                # 保存HTML文件
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(html_content)
                
                print(f"成功保存公众号链接: {filename} -> {file_path}")
                return True
            
            # 使用简化的请求头进行下载
            download_headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36 Edg/141.0.0.0',
                'Accept': '*/*',
                'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
                'Accept-Encoding': 'gzip, deflate, br',
                'Connection': 'keep-alive',
                'Referer': 'https://ima.qq.com/',
                'Sec-Fetch-Dest': 'document',
                'Sec-Fetch-Mode': 'navigate',
                'Sec-Fetch-Site': 'cross-site',
                'Sec-Fetch-User': '?1'
            }
            
            # 获取文档内容
            response = self.session.get(jump_url, headers=download_headers, timeout=30)
            response.raise_for_status()
            
            # 尝试从Content-Disposition头获取文件名
            content_disposition = response.headers.get('Content-Disposition', '')
            if 'filename=' in content_disposition:
                filename = content_disposition.split('filename=')[1].strip('"')
                filename = self.sanitize_filename(filename)
            else:
                # 根据内容类型确定文件扩展名
                content_type = response.headers.get('Content-Type', '')
                if 'pdf' in content_type:
                    filename = f"{safe_title}.pdf"
                elif 'word' in content_type or 'document' in content_type:
                    filename = f"{safe_title}.docx"
                elif 'text' in content_type:
                    filename = f"{safe_title}.txt"
                else:
                    # 尝试从URL获取扩展名
                    parsed_url = urlparse(jump_url)
                    path = parsed_url.path
                    if '.' in path:
                        ext = path.split('.')[-1]
                        filename = f"{safe_title}.{ext}"
                    else:
                        filename = f"{safe_title}.html"
            
            # 检查并修复文件名中的重复后缀
            filename = self.fix_duplicate_extension(filename)
            
            # 重新检查文件是否已存在（因为实际文件名可能与预期不同）
            file_path = os.path.join(save_dir, filename)
            if os.path.exists(file_path):
                print(f"文件已存在，跳过下载: {filename}")
                return True
            
            # 保存文件
            with open(file_path, 'wb') as f:
                f.write(response.content)
            
            print(f"成功下载: {filename} -> {file_path}")
            return True
            
        except requests.exceptions.RequestException as e:
            print(f"下载失败 {title}: {e}")
            # 尝试使用更简单的请求头重试
            try:
                print(f"尝试使用简单请求头重新下载: {title}")
                simple_headers = {
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
                }
                response = self.session.get(jump_url, headers=simple_headers, timeout=30)
                response.raise_for_status()
                
                # 检查并修复文件名中的重复后缀
                retry_filename = f"{safe_title}.pdf"
                retry_filename = self.fix_duplicate_extension(retry_filename)
                
                # 检查文件是否已存在
                file_path = os.path.join(save_dir, retry_filename)
                if os.path.exists(file_path):
                    print(f"文件已存在，跳过重试下载: {retry_filename}")
                    return True
                
                # 保存文件
                with open(file_path, 'wb') as f:
                    f.write(response.content)
                
                print(f"成功下载: {retry_filename} -> {file_path}")
                return True
                
            except Exception as retry_e:
                print(f"重试下载也失败 {title}: {retry_e}")
                return False
        except Exception as e:
            print(f"保存文件失败 {title}: {e}")
            return False
    
    def get_main_folder_name(self, knowledge_data):
        """
        从API响应中获取主目录名称
        
        Args:
            knowledge_data: 知识列表数据
            
        Returns:
            str: 主目录名称
        """
        # 优先从knowledge_base_info获取主目录名称
        if 'knowledge_base_info' in knowledge_data and 'basic_info' in knowledge_data['knowledge_base_info']:
            main_name = knowledge_data['knowledge_base_info']['basic_info'].get('name', '')
            if main_name:
                return self.sanitize_filename(main_name)
        
        # 如果knowledge_base_info中没有，从current_path获取
        current_path = knowledge_data.get('current_path', [])
        if current_path and len(current_path) > 0:
            # 获取最后一个路径项的名称（通常是当前目录）
            current_item = current_path[-1]
            if 'name' in current_item:
                return self.sanitize_filename(current_item['name'])
        
        return "未知目录"
    
    def get_current_path_name(self, knowledge_data):
        """
        从current_path或knowledge_base_info获取当前目录名称
        
        Args:
            knowledge_data: 知识列表数据
            
        Returns:
            str: 当前目录名称
        """
        # 优先从current_path获取名称
        current_path = knowledge_data.get('current_path', [])
        if current_path and len(current_path) > 0:
            # 获取最后一个路径项的名称（通常是当前目录）
            current_item = current_path[-1]
            if 'name' in current_item:
                return current_item['name']
        
        # 如果current_path为空，从knowledge_base_info获取
        if 'knowledge_base_info' in knowledge_data and 'basic_info' in knowledge_data['knowledge_base_info']:
            return knowledge_data['knowledge_base_info']['basic_info'].get('name', '')
        
        return "未知目录"
    
    def download_all_documents(self, knowledge_data, save_dir="downloads"):
        """
        使用多线程下载所有文档
        
        Args:
            knowledge_data: 知识列表数据
            save_dir: 保存目录
            
        Returns:
            dict: 下载结果统计
        """
        if not knowledge_data or 'knowledge_list' not in knowledge_data:
            print("无效的知识数据")
            return {"success": 0, "failed": 0, "total": 0}
        
        knowledge_list = knowledge_data.get('knowledge_list', [])
        if not knowledge_list:
            print("知识列表为空")
            return {"success": 0, "failed": 0, "total": 0}
        
        # 从API响应中获取主目录名称
        main_folder_name = self.get_main_folder_name(knowledge_data)
        
        # 创建主目录结构：downloads/主目录名称/
        main_save_dir = os.path.join(save_dir, main_folder_name)
        os.makedirs(main_save_dir, exist_ok=True)
        print(f"创建主目录: {os.path.abspath(main_save_dir)}")
        
        print(f"开始多线程下载 {len(knowledge_list)} 个文档到目录: {os.path.abspath(main_save_dir)}")
        print(f"使用 {MAX_WORKERS} 个线程进行下载")
        
        # 重置计数器
        self.success_count = 0
        self.failed_count = 0
        
        # 使用线程池进行多线程下载
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            # 提交所有下载任务
            future_to_item = {}
            for i, item in enumerate(knowledge_list, 1):
                future = executor.submit(
                    self.download_single_document, 
                    item, 
                    main_save_dir,  # 使用主目录路径
                    i, 
                    len(knowledge_list)
                )
                future_to_item[future] = item
            
            # 处理完成的任务
            completed_count = 0
            for future in as_completed(future_to_item):
                completed_count += 1
                try:
                    result = future.result()
                    if result["success"]:
                        if "folder_result" in result:
                            # 文件夹下载结果
                            folder_result = result["folder_result"]
                            print(f"✓ 文件夹完成: {result['title']} (包含 {folder_result['success']} 个文件)")
                        else:
                            print(f"✓ 文件完成: {result['title']}")
                    else:
                        print(f"✗ 失败: {result['title']} - {result.get('reason', '未知错误')}")
                except Exception as e:
                    print(f"✗ 异常: {future_to_item[future].get('title', '未知')} - {e}")
                    with self.lock:
                        self.failed_count += 1
                
                # 显示进度
                if completed_count % 10 == 0 or completed_count == len(knowledge_list):
                    with self.lock:
                        print(f"进度: {completed_count}/{len(knowledge_list)} "
                              f"(成功: {self.success_count}, 失败: {self.failed_count})")
                
                # 添加小延迟避免请求过于频繁
                time.sleep(DOWNLOAD_DELAY)
        
        result = {
            "success": self.success_count,
            "failed": self.failed_count,
            "total": len(knowledge_list)
        }
        
        print(f"\n多线程下载完成! 成功: {self.success_count}, 失败: {self.failed_count}, 总计: {len(knowledge_list)}")
        return result


def main():
    """主函数示例"""
    downloader = IMAKnowledgeDownloader()
    
    # 获取所有页面的知识列表
    print("正在获取知识列表...")
    knowledge_data = downloader.get_all_knowledge_list()
    
    if knowledge_data and knowledge_data.get('code') == 0:
        # 下载所有文档到ima目录下的downloads文件夹
        result = downloader.download_all_documents(knowledge_data, "downloads")
        
        # 打印结果
        print(f"\n下载统计:")
        print(f"成功: {result['success']}")
        print(f"失败: {result['failed']}")
        print(f"总计: {result['total']}")
    else:
        print("获取知识列表失败")
        print("\n提示：如果API请求失败，可能是因为：")
        print("1. 需要登录或认证")
        print("2. API端点已更改")
        print("3. 请求参数不正确")
        print("4. 网络连接问题")


if __name__ == "__main__":
    main()
